# dao/postgresql_connect.py
"""
连接 PostgreSQL 的配置
"""

import os
from dotenv import load_dotenv
from sqlalchemy import create_engine, text
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

load_dotenv()

postgres_host = os.getenv("POSTGRES_HOST")
postgres_port = os.getenv("POSTGRES_PORT")
postgres_db = os.getenv("POSTGRES_DB")
postgres_user = os.getenv("POSTGRES_USER")
postgres_password = os.getenv("POSTGRES_PASSWORD")

# 添加时区设置
DATABASE_URL = f"postgresql://{postgres_user}:{postgres_password}@{postgres_host}:{postgres_port}/{postgres_db}"


def check_connection() -> bool:
    """
    检查PostgreSQL连接是否正常
    如果连接失败会抛出异常
    """
    try:
        engine = create_engine(DATABASE_URL)
        with engine.connect() as connection:
            # 设置时区
            connection.execute(text("SET timezone TO 'Asia/Shanghai'"))
            connection.execute(text("SELECT 1"))
            print("PostgreSQL连接成功!")
            return True
    except Exception as e:
        print(f"PostgreSQL连接错误: {e}")
        raise Exception(f"PostgreSQL连接检查失败: {str(e)}")


# 创建数据库引擎和会话
engine = create_engine(
    DATABASE_URL,
    connect_args={
        "options": "-c timezone=Asia/Shanghai"
    }
)

# 创建会话工厂，确保每个会话都使用正确的时区
def get_db():
    db = SessionLocal()
    try:
        # 设置会话时区
        db.execute(text("SET timezone TO 'Asia/Shanghai'"))
        yield db
    finally:
        db.close()

SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()
